
from manim import *
from pathlib import Path
import os

FLAGS = f"-qh"
SCENE = "Scene"

Text.set_default(font_size = 24)
MathTex.set_default(font_size = 36) 

COLORS = {
    "-1": BLUE_E,
    "0" : BLUE_D,
    "1" : BLUE_B,
    "single": GOLD,
    "result": RED,
    "row": TEAL,
    "col": GREEN
}

class Scene(Scene):

    def construct(self):
        like_positions = []
        dislike_positions = []
        neutral_positions = []

        ########## Scene1 ##########

        t0 = Table(
            [["喜欢", "喜欢", "不喜欢", "喜欢"],
            ["喜欢","无感", "不喜欢", "喜欢"],
            ["不喜欢", "无感", "喜欢", "无感"],
            ["喜欢", "不喜欢", "无感", "喜欢"]],
            row_labels=[Text("A"), Text("B"), Text("C"), Text("D")],
            col_labels=[Text("让子弹飞"), Text("霸王别姬"), Text("喜剧之王"), Text("驴得水")],
            ).scale(0.7).shift(LEFT * 2.5 + UP * 1.0)

        # 获取颜色位置
        for i in range(1, 5, 1):
            for j in range(1, 5, 1): 
                if str(t0.get_rows()[i][j].lines_text).split("'")[1] == "不喜欢":
                    dislike_positions.append((i, j))
                elif str(t0.get_rows()[i][j].lines_text).split("'")[1] == "无感":
                    neutral_positions.append((i, j))
                else:
                    like_positions.append((i, j))

        ########## Scene2 ##########

        t1 = Table(
            [["1", "1", "-1", "1"],
            ["1","0", "-1", "1"],
            ["-1", "0", "1", "0"],
            ["1", "-1", "0", "1"]],
            row_labels=[Text("A"), Text("B"), Text("C"), Text("D")],
            col_labels=[Text("让子弹飞"), Text("霸王别姬"), Text("末代皇帝"), Text("驴得水")],
            ).scale(0.7).shift(LEFT * 2.5 + UP * 1.0)
        
        # 加颜色
        t0.get_rows()[0].set_color(COLORS["row"])
        t0.get_columns()[0].set_color(COLORS["col"])
        t1.get_rows()[0].set_color(COLORS["row"])
        t1.get_columns()[0].set_color(COLORS["col"])

        for pos in dislike_positions:
            t0.get_rows()[pos[0]][pos[1]].set_color(COLORS["-1"])
            t1.get_rows()[pos[0]][pos[1]].set_color(COLORS["-1"])

        for pos in neutral_positions:
            t0.get_rows()[pos[0]][pos[1]].set_color(COLORS["0"])
            t1.get_rows()[pos[0]][pos[1]].set_color(COLORS["0"])

        for pos in like_positions:
            t0.get_rows()[pos[0]][pos[1]].set_color(COLORS["1"])
            t1.get_rows()[pos[0]][pos[1]].set_color(COLORS["1"])



        ########## Scene3 ##########


        line0 = MathTex(r"(-1)", r"\times", r"(-1)", r"=", r"-1")
        line1 = MathTex(r"1"   , r"\times", r"1"   , r"=", r"1"   )
        line2 = MathTex(r"(-1)", r"\times", r"(-1)", r"=", r"1"   )
        line3 = MathTex(r"1"   , r"\times", r"(-1)", r"=", r"-1")
        line4 = MathTex(r"0"   , r"\times", r"1"   , r"=", r"0"   )


        for line in [line1, line2, line3, line4]:
            for i in [0, 2]:
                if str(line[i]).split("'")[1] == "(-1)":
                    line[i].set_color(COLORS["-1"])
                elif str(line[i]).split("'")[1] == "0":
                    line[i].set_color(COLORS["0"])
                else:
                    line[i].set_color(COLORS["1"])
            line[-1].set_color(COLORS["single"])

        text0 = Text("两字").set_color(COLORS["single"])
        text1 = Text("相同").set_color(COLORS["single"])
        text2 = Text("相同").set_color(COLORS["single"])
        text3 = Text("不同").set_color(COLORS["single"])
        text4 = Text("无关").set_color(COLORS["single"])

        
        Y_OFFSET = 3.2
        X_OFFSET = 3.6
        for i in range(len(line0)):
            line0[i].shift(RIGHT * (i * 0.1 + X_OFFSET) + UP * Y_OFFSET)
            line1[i].move_to(line0[i].get_center()).shift(DOWN * (1 + 0 * 0.7))
            line2[i].move_to(line0[i].get_center()).shift(DOWN * (1 + 1 * 0.7))
            line3[i].move_to(line0[i].get_center()).shift(DOWN * (1 + 2 * 0.7))
            line4[i].move_to(line0[i].get_center()).shift(DOWN * (1 + 3 * 0.7))
        

        text0.next_to(line0[-1].get_center(), RIGHT * 2.5)
        text1.next_to(line1[-1].get_center(), RIGHT * 2.5)
        text2.next_to(line2[-1].get_center(), RIGHT * 2.5)
        text3.next_to(line3[-1].get_center(), RIGHT * 2.5)
        text4.next_to(line4[-1].get_center(), RIGHT * 2.5)


        ########## Scene3.5 ##########
        dot_product = MathTex(r"\mathbf{u} \cdot \mathbf{v} = \sum_{k} u[k] v[k]")
        dot_product.set_color(COLORS["result"])
        dot_product.shift(DOWN * 2.0)

        ########## Scene4 ##########
        eq1 = MathTex(
        r"\mathbf{A}", r"\cdot", r"\mathbf{B}", r"=", 
        r"(", r"1" , r"\times", r"1" , r")", r"+", 
        r"(", r"1" , r"\times", r"0" , r")", r"+",
        r"(", r"-1" , r"\times", r"-1", r")", r"+",
        r"(", r"1", r"\times", r"1" , r")",
        r"=", r"3"
        )
        eq2 = MathTex(
        r"\mathbf{A}", r"\cdot", r"\mathbf{C}", r"=", 
        r"(", r"1" , r"\times", r"-1" , r")", r"+", 
        r"(", r"1" , r"\times", r"0"  , r")", r"+",
        r"(", r"-1" , r"\times", r"1"  , r")", r"+",
        r"(", r"1", r"\times", r"0"  , r")",
        r"=", r"-2"
        )

        eq1[0].set_color(COLORS["col"])
        eq1[2].set_color(COLORS["col"])
        eq2[0].set_color(COLORS["col"])
        eq2[2].set_color(COLORS["col"])

        for eq in [eq1, eq2]:
            for i in [5, 7, 11, 13, 17, 19, 23, 25]:
                if str(eq[i]).split("'")[1] == "-1":
                    eq[i].set_color(COLORS["-1"])
                elif str(eq[i]).split("'")[1] == "0":
                    eq[i].set_color(COLORS["0"])
                else:
                    eq[i].set_color(COLORS["1"])
            eq[-1].set_color(COLORS["result"])

        eq3 = MathTex(r"3", r">", r"-2", font_size=45)
        # eq3[0][0].set_color(COLORS["result"])
        # eq3[0][2].set_color(COLORS["result"])
        eq3.set_color(COLORS["result"])

        eq1.shift(DOWN * 1.5 + LEFT * 1.5)
        eq2.shift(DOWN * 2.5 + LEFT * 1.5)
        eq3.shift(DOWN * 2.0 + RIGHT * 4)








        ########## Animation ##########

        ########## Scene1 ##########

        self.wait(1)
        self.play(Write(t0), run_time=3)
        self.wait(1)



        ########## Scene2 ##########
        # self.add(t0)

        self.wait(5)
        self.play(
            *[ReplacementTransform(t0.get_rows()[pos[0]][pos[1]], t1.get_rows()[pos[0]][pos[1]]) for pos in like_positions]
        )
        self.wait(1)
        self.wait(1)
        self.play(
            *[ReplacementTransform(t0.get_rows()[pos[0]][pos[1]], t1.get_rows()[pos[0]][pos[1]]) for pos in neutral_positions]
        )
        self.wait(1)
        self.wait(1)
        self.play(
            *[ReplacementTransform(t0.get_rows()[pos[0]][pos[1]], t1.get_rows()[pos[0]][pos[1]]) for pos in dislike_positions]
        )
        self.wait(1)


        ########## Scene3 ##########
        # self.add(t1)

        one_one = VGroup(t1.get_rows()[1][1], t1.get_rows()[2][1])
        neg_neg = VGroup(t1.get_rows()[1][3], t1.get_rows()[2][3])
        one_neg = VGroup(t1.get_rows()[1][1], t1.get_rows()[3][1])
        one_zer = VGroup(t1.get_rows()[1][2], t1.get_rows()[2][2])

        self.wait()
        self.play(Indicate(one_one), run_time=1.5)
        self.wait()
        self.play(ReplacementTransform(one_one.copy(), line1))
        self.wait()
        self.play(FadeIn(text1))

        self.wait()
        self.play(Indicate(neg_neg), run_time=1.5)
        self.wait()
        self.play(ReplacementTransform(neg_neg.copy(), line2))
        self.wait()
        self.play(FadeIn(text2))

        self.wait()
        self.play(Indicate(one_neg), run_time=1.5)
        self.wait()
        self.play(ReplacementTransform(one_neg.copy(), line3))
        self.wait()
        self.play(FadeIn(text3))

        self.wait()
        self.play(Indicate(one_zer), run_time=1.5)
        self.wait()
        self.play(ReplacementTransform(one_zer.copy(), line4))
        self.wait()
        self.play(FadeIn(text4))

        ########## Scene3.5 ##########
        self.wait()
        self.play(FadeIn(dot_product), run_time=2)
        self.wait()
        self.play(Circumscribe(dot_product))
        self.wait(4)
        self.play(FadeOut(dot_product), run_time=2)
        self.wait()

        ########## Scene4 ##########
        # self.add(t1)
        # self.add(line1, text1, line2, text2, line3, text3, line4, text4)

        self.wait(5)
        self.play(
            TransformMatchingShapes(t1.get_rows()[1][0].copy(), eq1[0]),
            TransformMatchingShapes(t1.get_rows()[1][1].copy(), eq1[5]),
            TransformMatchingShapes(t1.get_rows()[1][2].copy(), eq1[11]),
            TransformMatchingShapes(t1.get_rows()[1][3].copy(), eq1[17]),
            TransformMatchingShapes(t1.get_rows()[1][4].copy(), eq1[23]),
            TransformMatchingShapes(t1.get_rows()[2][0].copy(), eq1[2]),
            TransformMatchingShapes(t1.get_rows()[2][1].copy(), eq1[7]),
            TransformMatchingShapes(t1.get_rows()[2][2].copy(), eq1[13]),
            TransformMatchingShapes(t1.get_rows()[2][3].copy(), eq1[19]),
            TransformMatchingShapes(t1.get_rows()[2][4].copy(), eq1[25]),
            FadeIn(eq1[1], eq1[3], eq1[4], eq1[6], eq1[8], eq1[9], eq1[10], eq1[12], eq1[14], eq1[15], eq1[16], eq1[18], eq1[20], eq1[21], eq1[22], eq1[24], eq1[26], eq1[27], eq1[28]),
            run_time=1.5
        )
        self.wait()

        self.wait()
        self.play(
            TransformMatchingShapes(t1.get_rows()[1][0].copy(), eq2[0]),
            TransformMatchingShapes(t1.get_rows()[1][1].copy(), eq2[5]),
            TransformMatchingShapes(t1.get_rows()[1][2].copy(), eq2[11]),
            TransformMatchingShapes(t1.get_rows()[1][3].copy(), eq2[17]),
            TransformMatchingShapes(t1.get_rows()[1][4].copy(), eq2[23]),
            TransformMatchingShapes(t1.get_rows()[3][0].copy(), eq2[2]),
            TransformMatchingShapes(t1.get_rows()[3][1].copy(), eq2[7]),
            TransformMatchingShapes(t1.get_rows()[3][2].copy(), eq2[13]),
            TransformMatchingShapes(t1.get_rows()[3][3].copy(), eq2[19]),
            TransformMatchingShapes(t1.get_rows()[3][4].copy(), eq2[25]),
            FadeIn(eq2[1], eq2[3], eq2[4], eq2[6], eq2[8], eq2[9], eq2[10], eq2[12], eq2[14], eq2[15], eq2[16], eq2[18], eq2[20], eq2[21], eq2[22], eq2[24], eq2[26], eq2[27], eq2[28]),
            run_time=1.5
        )
        self.wait()


        ########## Scene5 ##########
        # self.add(t1)
        # self.add(line1, text1, line2, text2, line3, text3, line4, text4)
        # self.add(eq1)
        # self.add(eq2)

        self.wait(4)
        self.play(
            TransformMatchingShapes(
                Group(eq1[-1], eq2[-1]).copy(),
                eq3
            ), run_time=1.5
        )
        self.wait()

        ########## Scene5 ##########
        self.wait()
        self.play(Wiggle(eq3[0], scale_value=1.3), run_time=2)
        self.wait()
        self.play(Wiggle(eq3[-1], scale_value=1.3), run_time=2)
        self.wait()







if __name__  == '__main__':
    script_name = f"{Path(__file__).resolve()}"
    os.system(f"manim {FLAGS} {script_name} {SCENE}")